import matplotlib.pyplot as plt
import numpy as np
"""
#洛伦兹方程
Dx/dt = - sigma * (X - Y)
DY/dt = gama*X - Y -X*Z 
DZ/dt = - beta*Z + X*Y
"""
#指定参数
sigma = 10
gamma = 25
beta = 8/3
t = [0]
x = [1]
y = [0]
z = [0]
h = 0.01
N = 10000

def Lorenz(x0,y0,z0,p,q,r,T):
  #微分迭代步长
  h=0.01
  x=[]
  y=[]
  z=[]
  for t in range(T):
    xt=x0+h*p*(y0-x0)
    yt=y0+h*(q*x0-y0-x0*z0)
    zt=z0+h*(x0*y0-r*z0)
    #x0、y0、z0统一更新
    x0,y0,z0=xt,yt,zt
    x.append(x0)
    y.append(y0)
    z.append(z0)
  return x,y,z

#Dx/dt = - sigma * (X - Y)
def FX(X, Y, Z):
    return - sigma * (X - Y)

#DY/dt = gama*X - Y -X*Z
def FY(X, Y, Z):
    return gamma * X - Y - X * Z

#DZ/dt = - beta*Z + X*Y
def FZ(X, Y, Z):
    return - beta * Z + X * Y

#龙格库塔方程组
def main():
    for i in range(N):
        K1 = FX(x[-1], y[-1], z[-1])
        L1 = FY(x[-1], y[-1], z[-1])
        M1 = FZ(x[-1], y[-1], z[-1])
        K2 = FX(x[-1] + h * K1 / 2, y[-1] + h * L1 / 2, z[-1] + h * M1 / 2)
        L2 = FY(x[-1] + h * K1 / 2, y[-1] + h * L1 / 2, z[-1] + h * M1 / 2)
        M2 = FZ(x[-1] + h * K1 / 2, y[-1] + h * L1 / 2, z[-1] + h * M1 / 2)
        K3 = FX(x[-1] + h * K2 / 2, y[-1] + h * L2 / 2, z[-1] + h * M2 / 2)
        L3 = FY(x[-1] + h * K2 / 2, y[-1] + h * L2 / 2, z[-1] + h * M2 / 2)
        M3 = FZ(x[-1] + h * K2 / 2, y[-1] + h * L2 / 2, z[-1] + h * M2 / 2)
        K4 = FX(x[-1] + h * K3, y[-1] + h * L3, z[-1] + h * M3)
        L4 = FY(x[-1] + h * K3, y[-1] + h * L3, z[-1] + h * M3)
        M4 = FZ(x[-1] + h * K3, y[-1] + h * L3, z[-1] + h * M3)
        x.append(x[-1] + h / 6 * (K1 + 2 * K2 + 2 * K3 + K4))
        y.append(y[-1] + h / 6 * (L1 + 2 * L2 + 2 * L3 + L4))
        z.append(z[-1] + h / 6 * (M1 + 2 * M2 + 2 * M3 + M4))
    #画三维图
    fig = plt.figure()
    ax = fig.add_subplot(projection='3d')
    ax.plot(x, y, z, label='parametric curve')
    plt.show()

    # 设初值
    x0 = -16
    y0 = -21
    z0 = 33
    x1, y1, z1 = Lorenz(x0, y0, z0, sigma, gamma, beta, N)
    # 初值微小的变化
    x0 = -16
    y0 = -21
    z0 = 33.00001
    xx, yy, zz = Lorenz(x0, y0, z0, sigma, gamma, beta, N)
    t = np.arange(0, N)
    plt.scatter(t, x1, s=1)
    plt.scatter(t, xx, s=1)
    plt.show()

if __name__=='__main__':
  main()
